"""--------------------------------------------------------------------------------------------------
Filename:   liveplot.py
Author:     Ben Han (uha10@uclive.ac.nz)
Date:       17 May 2010
Description:
    Main plotting module for biofeedback application.
    Contains the GUI layer and plotting functions.
--------------------------------------------------------------------------------------------------"""

import matplotlib
# The recommended way to use wx with mpl is with the WXAgg backend. 
matplotlib.interactive(False)
matplotlib.use('WXAgg')

import numpy as np
import pylab
import wx, traceback
import database
from targetfactory import TargetFactory
from decimal import Decimal
from plotpanel import PlotControl, PlotPanel
import threading
import gc
import datetime


#===============================================================================
# CLASS: LivePlot(wx.Frame)
#    Real-time data plotting class.
#    This class owns the data arrays containing the EMG and time values for
#    the plot points.
#===============================================================================
class LivePlot(wx.Frame):
    def __init__(self, app, parent=None, id=-1, title="Live plot frame"):
        wx.Frame.__init__(self, parent, id, title, size=(800, 600))
        self.app = app
        self.top_panel = wx.Panel(self, -1)        # top-level panel for frame
        self.plotpanel = PlotPanel(parent=self.top_panel, color=None, dpi=80, id=-1)
        self.plotctrl = PlotControl(parent=self.top_panel, id=-1, app=app)

        self.__do_layout()
        self.__setup_plot()
        
    def __do_layout(self):
        sizer_frame = wx.BoxSizer(wx.VERTICAL)  # sizer for frame
        sizer_main = wx.BoxSizer(wx.VERTICAL)
        sizer_main.Add(self.plotpanel, 1, wx.EXPAND | wx.ALL, 0)
        sizer_main.Add(self.plotctrl, 0, wx.EXPAND | wx.ALL, 5)
        self.top_panel.SetSizer(sizer_main)
        sizer_frame.Add(self.top_panel, 1, wx.EXPAND, 0)
        self.SetSizer(sizer_frame)
        sizer_frame.Fit(self)
        self.ShowFullScreen(not self.IsFullScreen(), wx.FULLSCREEN_ALL)
        self.Layout()
        
    def __setup_plot(self):
        # Instantiate a matplotlib axes object
        self.axes = self.plotpanel.figure.add_subplot(111)    # self.axes is an instance of an Axes object
        self.axes.set_axis_bgcolor(self.app.cfg._plot_back_color)
        self.axes.set_title('sEMG level', size=20)
        self.axes.set_autoscale_on(False)

        pylab.setp(self.axes.get_xticklabels(), fontsize=13)
        pylab.setp(self.axes.get_yticklabels(), fontsize=13)

        # Instantiate a target factory to be used in skill training
        self.tfactory = TargetFactory(self.axes, self.app.cfg, self.app.sess_status)
        
        # Initialise data arrays to be added to later
        self.app.data_lock.acquire()    # ACQUIRE DATA_LOCK <--
        self.x_data = np.array([])      # array of x values for the plot
        self.y_data = np.array([])      # array of y values for the plot
        self.emgdata = []               # used to store raw EMG data
        self.plot_state = []            # used to store state information about the plot for every time point
        self.app.data_lock.release()    # RELEASE DATA_LOCK -->
        
        self.sample_period = Decimal(1) / Decimal(str(self.app.cfg._sample_rate))
        
        # Plot the data and grab a reference to the line added
        self.emgplot = self.axes.plot(self.x_data,
                                      self.y_data,
                                      linewidth=self.app.cfg._plot_line_width,
                                      color=self.app.cfg._plot_line_color,
                                      animated=True)[0]
                
        self.init_plot()
        
    def __remove_target(self):
        # Remove old target
        if self.app.sess_status.target_on:
            self.target.remove()
            self.app.sess_status.target_on = False

    def __draw_cursors(self):
        line_color = 'black'
        line_width = 2
        if self.app.ctrl.chkbx_ctrl.cb_time.IsChecked():
            try:
                if self.app.sess_status.timeline:
                    self.time_line.remove()
                t = float(self.app.sess_status.y_index) - float(self.sample_period)
                self.time_line = self.axes.axvline(t, color=line_color, linewidth=line_width)
                self.app.sess_status.timeline = True
            except:
                print traceback.print_exc()
        else:
            if self.app.sess_status.timeline:
                self.time_line.remove()
                self.app.sess_status.timeline = False

        if self.app.ctrl.chkbx_ctrl.cb_ampl.IsChecked():
            try:
                if self.app.sess_status.amplline:
                    self.ampl_line.remove()
                if len(self.y_data) > 0:
                    self.ampl_line = self.axes.axhline(self.y_data[-1], color=line_color, linewidth=line_width)
                    self.app.sess_status.amplline = True
            except:
                print traceback.format_exc(10)
        else:
            if self.app.sess_status.amplline:
                self.ampl_line.remove()
                self.app.sess_status.amplline = False
        
    def __draw_target_strength(self, init=False):
        self.__remove_target()
        adj = 1 - (float(self.app.cfg._str_increment) / 100)
        
        if init:
            self.app.sess_status.str_target = self.app.sess_status.baseline
            self.app.ctrl.ymax_ctrl.set_value(self.app.sess_status.baseline * 1.8)    # Scale the Y-axis range
            self.app.cfg._y_axis_max = int(self.app.ctrl.ymax_ctrl.get_value())
            self.app.sess_status.peak = (0,0)
            self.app.sess_status.result = False
        elif self.app.cfg._str_running_count:
            # Check the count and adjust target as necessary
            if self.app.sess_status.count >= self.app.cfg._str_hit_tol:      # if target has been hit consecutively, decrease size
                self.app.sess_status.str_target *= adj    # decrease size
                self.app.sess_status.count = 0
            elif self.app.sess_status.count <= -(self.app.cfg._str_miss_tol):   # if target has been missed consecutively, increase size
                self.app.sess_status.str_target /= adj    # increase size
                self.app.sess_status.count = 0
        else:
            # Check the count and adjust target as necessary
            if self.app.sess_status.target_hits >= self.app.cfg._str_hit_tol:        # if target has been hit consecutively, decrease size
                self.app.sess_status.str_target /= adj    # increase size
                self.app.sess_status.target_hits = 0
            elif self.app.sess_status.target_misses >= self.app.cfg._str_miss_tol:   # if target has been missed consecutively, increase size
                self.app.sess_status.str_target *= adj    # decrease size
                self.app.sess_status.target_misses = 0
            
        # Draw line on plot
        color = self.app.cfg._target_color
        self.target = self.axes.axhline(self.app.sess_status.str_target, color=color, lw=self.app.cfg._target_line_width)
        self.app.sess_status.target_on = True
        
        # Print target details to console
        print "=== Target parameters ================================"
        print "    Amplitude:    %.2f" % self.app.sess_status.str_target
        print "======================================================"
        
        # reset the axes
        self.axes.set_xbound(lower=0, upper=self.app.cfg._trial_len)
        self.axes.set_ybound(lower=0, upper=self.app.cfg._y_axis_max)

    def __draw_target_skill(self, init=False):
        self.__remove_target()
        if init:
            self.target = self.tfactory.recalc(0, init)
            self.app.sess_status.result = False
        else:
            if self.app.cfg._running_count:
                # Check hit count and adjust target as necessary
                if self.app.sess_status.count >= self.app.cfg._hit_tol:         # if target has been hit consecutively, decrease size
                    self.target = self.tfactory.recalc(-1)
                    self.app.sess_status.count = 0
                elif self.app.sess_status.count <= -(self.app.cfg._miss_tol):   # if target has been missed consecutively, increase size
                    self.target = self.tfactory.recalc(1)
                    self.app.sess_status.count = 0
                else:
                    self.target = self.tfactory.recalc(0, init)
            else:
                if self.app.sess_status.target_hits >= self.app.cfg._hit_tol:        # if target has been hit consecutively, decrease size
                    self.target = self.tfactory.recalc(-1)
                    self.app.sess_status.target_hits = 0
                elif self.app.sess_status.target_misses >= self.app.cfg._miss_tol:   # if target has been missed consecutively, increase size
                    self.target = self.tfactory.recalc(1)
                    self.app.sess_status.target_misses = 0
                else:
                    self.target = self.tfactory.recalc(0, init)
        self.app.sess_status.target_on = True

        # reset the axes
        self.axes.set_xbound(lower=0, upper=self.app.cfg._trial_len)
        self.axes.set_ybound(lower=0, upper=self.app.cfg._y_axis_max)
        
        # set delay marker, get the time value of right-hand edge of target 
        self.app.sess_status.delay_mark = Decimal( str( float(self.tfactory.x2) + float(self.app.cfg._delay_len) ) )
        
    def update_data(self, emglevel):
        # This function is called by the serial thread to append EMG values to the data array.
        # Any work to be done by the serial thread from this function should 
        # be added to the app.job_queue object rather than performed directly. 
        s_state = self.app.sess_status
        jobs = self.app.job_queue
        if s_state.alive:
            try:
                self.app.data_lock.acquire()        # ACQUIRE DATA_LOCK <--
                if s_state.y_index <= self.app.cfg._trial_len:
                    # Test the latest value in the data array and update maximum if needed
                    if emglevel > s_state.peak[1]:
                        s_state.peak = (s_state.y_index, emglevel)
        
                    # Add the latest value to the data array and increment y_index
                    self.y_data = np.append(self.y_data, emglevel)
                    self.x_data = np.append(self.x_data, s_state.y_index)
                    s_state.y_index += self.sample_period
                    # Store plot state information
                    self.plot_state.append(self.emgplot.get_visible())
                    
                    # Add a draw() call to the job queue
                    jobs.put(self.draw)
                    
                else:
                    # End of plot
                    s_state.alive = False           # turn itself off so end_trial is only added to the queue once
                    if s_state.mode == 2 or s_state.mode == 4:
                        jobs.put(self.end_trial)    # end the trial
                    else:
                        jobs.put(self.init_plot)    # calibration mode so just reset
            except:
                print "Exception in update_data()"
                print traceback.format_exc()
            finally:
                self.app.data_lock.release()        # RELEASE DATA_LOCK -->
            
    def draw_grid(self, boo):
        # Turn the plot grid on or off
        if boo:
            self.axes.grid(True, color='gray')
        else:
            self.axes.grid(False)
        
    def status_msg(self, msg):
        # Set the status text control value
        self.plotctrl.set_status(msg)
        
    def end_trial(self):
        """ Resets the plot between trials.
            This function should only be called after a trial run has been completed
        """
        self.app.sess_status.alive = False      # stop the plot
        gc.collect()
        
        ti = self.app.sess_status.trial_index
        mode = self.app.sess_status.mode
        
        # Record the result of the last trial
        self.app.sess_status.result = self.test_target()
         
        # ----------------------- Strength trial
        if mode == 2:
            self.app.sess_status.type = "Strength"
            if self.app.sess_status.result:
                self.app.sess_status.target_hits += 1
                self.app.sess_status.target_misses = 0
                self.app.sess_status.count += self.app.cfg._hit_incr
            else:
                self.app.sess_status.target_misses += 1
                self.app.sess_status.target_hits = 0
                self.app.sess_status.count -= self.app.cfg._miss_incr

            # Scale Y axis if last trial went out of range
            if self.app.sess_status.ymax_auto and self.app.sess_status.peak[1] > (self.app.cfg._y_axis_max * 0.8):
                self.app.cfg._y_axis_max = self.app.sess_status.peak[1] * (1.2 + (self.app.cfg.y_margin / 100.0))

        # ----------------------- Skill trial
        elif mode == 4:
            self.app.sess_status.type = "Skill"
            if self.app.sess_status.result:
                self.app.sess_status.target_hits += 1
                self.app.sess_status.target_misses = 0
                self.app.sess_status.count += self.app.cfg._hit_incr
            else:
                self.app.sess_status.target_misses += 1
                self.app.sess_status.target_hits = 0
                self.app.sess_status.count -= self.app.cfg._miss_incr

        tr = self.get_trial()       # create a database.Trial object
        self.app.sess_lock.acquire()    # ACQUIRE SESS_LOCK <--
        self.app.status.session.add_trial(tr)
        self.app.sess_lock.release()    # RELEASE SESS_LOCK -->

        if ti > 0:
            # Save trial outcome to history list
            htuple = (self.app.sess_status.peak, tr.target, self.app.sess_status.result)
            self.app.sess_status.history.append(htuple)
        
        # Print a short summary of the trial to the command line
        if mode == 2:
            print ">> End   strength trial %i" % ti
        elif mode == 4:
            print ">> End   skill trial %i" % ti
        if self.app.cfg._str_running_count:
            print '> Cumulative  count: %i' % (self.app.sess_status.count)
            print '> Tolerances   Hits: %i\tMisses: %i\n' % (self.app.cfg._str_hit_tol, self.app.cfg._str_miss_tol)
        else:
            print '> Consecutive  Hits: %i/%i\tMisses: %i/%i' % (self.app.sess_status.target_hits, 
                                                                 self.app.cfg._str_hit_tol, 
                                                                 self.app.sess_status.target_misses, 
                                                                 self.app.cfg._str_miss_tol)
        print tr.get_summary()      # Print a short summary of the last trial to the console

        if ti % self.app.cfg._trials_per_block == 0:  # end block
            gc.collect()
            self.app.end_block()
            self.app.sess_status.block_index += 1
            # Wait until break is over (in worker thread)
            self.app.break_finished.wait()
            
        if ti >= self.app.cfg._trials_per_session:      # If all trials have been completed then end session
            gc.collect()
            self.app.end_session()
        else:
            # Initialise plot for the next run
            self.app.sess_status.trial_index += 1
            self.init_plot()
    
    def get_trial(self):
        # Returns a database.Trial object constructed from the current plot data
        
        # check that the array are the same length and truncate the longer if not
        lenx = len(self.x_data)
        leny = len(self.y_data)
        if not lenx == leny:
            diff = leny - lenx
            if diff > 0:                # if y_data has more items than x_data...
                self.y_data = self.y_data[:-diff]   # remove items from y_data to even up the arrays
            elif diff < 0:              # if x_data has more items than y_data...
                self.x_data = self.x_data[:-diff]   # remove items from x_data to even up the arrays
            
        # Record the raw emgdata
        for i in range(len(self.y_data)):
            self.emgdata.append((self.x_data[i], self.y_data[i], self.plot_state[i]))
            
        # Calculate axes size in display pixels
        xb = self.axes.get_xbound()
        yb = self.axes.get_ybound()
        tr = self.axes.transData
        ll = tr.transform((xb[0],yb[0]))
        ur = tr.transform((xb[1],yb[1]))
        psize = (ur[0]-ll[0], ur[1]-ll[1])
        sx = int(psize[0]+0.5)
        sy = int(psize[1]+0.5)
        
        # Create a trial object for this trial and append to the current session object
        tr = database.Trial(len=self.app.cfg._trial_len,
                            typ=self.app.sess_status.type,
                            peak=self.app.sess_status.peak,
                            mpl_obj=self.target, 
                            result=self.app.sess_status.result,
                            emgdata=self.emgdata,
                            plot_size=(sx, sy),
                            y_max=yb[1],
                            aspect=self.app.cfg._target_aspect_ratio)
        return tr

    def init_plot(self):
        """ Initial plot setup at the start of calibration and trials
        """
        self.reset_plot()
        
        alive = True       # begin adding jobs to the queue
        ci = self.app.sess_status.calib_index
        ti = self.app.sess_status.trial_index
        mode = self.app.sess_status.mode
        
        # ----------------------- Calibration modes
        if mode == 1:
            text = 'Strength Calibration: %i' % ci
            if ci == 0:
                alive = False
                self.app.ctrl.ymax_ctrl.set_auto(False)
                self.app.sess_status.ymax_auto = False

        elif mode == 3:
            text = 'Skill Calibration: %i' % ci
            if ci == 0:
                alive = False
                self.app.ctrl.ymax_ctrl.set_auto(False)
                self.app.sess_status.ymax_auto = False

        # ----------------------- Strength training mode
        elif mode == 2:
            ymax_baseline = self.app.cfg._y_margin_strength
            text = ('Strength trial: %i' % ti)
            if ti == 0:     # if pre-trial
                alive = False
                self.app.ctrl.ymax_ctrl.set_value(self.app.sess_status.baseline * ymax_baseline)    # Scale the Y-axis range
                self.app.cfg._y_axis_max = int(self.app.ctrl.ymax_ctrl.get_value())
                
            else:
                print "-----------------------------------------------------------------"
                print ">> Begin strength trial %i" % ti
                self.__draw_target_strength()

        # ----------------------- Skill training mode
        elif mode == 4:
            ymax_baseline = self.app.cfg._y_margin_skill
            text = ('Skill trial: %i' % ti)
            if ti == 0:     # if pre-trial
                alive = False     # disable plot update while waiting for beginning of trial 
                self.app.ctrl.ymax_ctrl.set_value(self.app.sess_status.baseline * ymax_baseline)    # Scale the Y-axis range
                self.app.cfg._y_axis_max = int(self.app.ctrl.ymax_ctrl.get_value())
                self.tfactory.set_size_limits()
            else:
                print "-----------------------------------------------------------------"
                print ">> Begin skill trial %i" % ti
                self.__draw_target_skill()      # draw skill target
                
                if self.app.cfg._use_fading:    # adjust visibility of plot if fading is being implemented
                    self.set_fading()
                
                if self.app.cfg._use_delay:     # adjust visibility of plot if delayed feedback is on
                    self.set_delay()
                
        else:
            text = ''
        self.status_msg(text)      # Update trial text
        
        # Set the serial data acquisition thread to write values to arrays
        self.app.sess_status.alive = alive
        
    def set_delay(self):
        # Used to implement delayed feedback
        # Turn off plot and update controls
        self.show_plot(False)
        # set time cursor on by default
        self.app.ctrl.chkbx_ctrl.cb_time.SetValue(True)
        self.plotctrl.cb_time.SetValue(True)
        
    def set_fading(self):
        block = self.app.sess_status.block_index
        sch = self.app.cfg.fading_schedule[block]
        on = self.app.sess_status.plots_on
        off = self.app.sess_status.plots_off
        phase = self.app.sess_status.fading_phase
        if phase:
            if off < sch[1]:
                self.__set_fading(True)
                self.app.sess_status.plots_off += 1    # increment number of trials that have used fading (i.e. plot not visible)
            else:
                self.__set_fading(False)
                self.app.sess_status.plots_on += 1
                self.app.sess_status.plots_off = 0          # reset counter
                self.app.sess_status.fading_phase = False   # change fading phase
        else:
            if on < sch[0]:
                self.__set_fading(False)
                self.app.sess_status.plots_on += 1     # increment number of trials with plot visible
            else:
                self.__set_fading(True)
                self.app.sess_status.plots_off += 1
                self.app.sess_status.plots_on = 0           # reset counter
                self.app.sess_status.fading_phase = True    # change fading phase
        
    def __set_fading(self, boo):
        # Used to implement fading schedule. Should be called at the start of each plot
        # init_plot() > set_fading() > __set_fading()
        boo = not boo       # invert value to set plot visible is equal to inverse of fading on
        self.show_plot(boo)
        
    def show_plot(self, boo):
        # set plot visibility and update controls
        self.app.ctrl.chkbx_ctrl.cb_plot.SetValue(boo)
        self.plotctrl.cb_plot.SetValue(boo)
        self.emgplot.set_visible(boo)
        
    def reset_plot(self):
        # Clears the plot frame and resets data arrays
        self.app.sess_status.peak = (0, 0)
        self.app.sess_status.result = False
        
        # reset data arrays
        self.app.data_lock.acquire()      # ACQUIRE DATA_LOCK <--
        self.app.sess_status.y_index = Decimal(0)
        self.x_data = np.array([])
        self.y_data = np.array([])
        self.emgdata = []
        self.plot_state = []
        self.emgplot.set_xdata(self.x_data)     # Set the plot data
        self.emgplot.set_ydata(self.y_data)
        self.app.data_lock.release()      # RELEASE DATA_LOCK -->

        # reset the axes
        self.axes.set_xbound(lower=0, upper=self.app.cfg._trial_len)
        self.axes.set_ybound(lower=0, upper=self.app.cfg._y_axis_max)
        
        # Set grid lines
        self.draw_grid(self.plotctrl.cb_grid.GetValue())

        # Update trial status text and redraw the plot
        self.app.display_status()
        try:
            canvas = self.plotpanel.canvas
            canvas.draw()
            self.background = canvas.copy_from_bbox(self.axes.bbox)
        except:
            print traceback.format_exc()
        self.Refresh()
        
    def draw(self):
        # Main draw function
        #    - Draws the changes in the model to the screen
        self.app.data_lock.acquire()      # ACQUIRE DATA_LOCK <--
        
        if self.app.sess_status == None:
            pass
        else:
            # Set y-axis limits
            if self.app.sess_status.ymax_auto:
                if self.app.sess_status.peak[1] > (self.app.cfg._y_axis_max * 0.8):
                    self.app.cfg._y_axis_max = self.app.sess_status.peak[1] * (1 + (self.app.cfg.y_margin / 100.0))
                    self.axes.set_ybound(lower=0, upper=self.app.cfg._y_axis_max)
    
            # Draw cursors if selected
            self.__draw_cursors()
            
            # Check if delayed feedback is being used, set plot visibility accordingly
            if self.app.sess_status.mode == 2 or self.app.sess_status.mode == 4:
                if self.app.cfg._use_delay:
                    delta = self.app.sess_status.y_index - self.app.sess_status.delay_mark 
                    if delta > 0:
                        self.show_plot(True)
            
            try:        # main canvas update function call
                # Restore the background
                self.plotpanel.canvas.restore_region(self.background)
                # Set the plot data
                self.emgplot.set_xdata(self.x_data)
                self.emgplot.set_ydata(self.y_data)
                # Update the emg plot
                self.axes.draw_artist(self.emgplot)
                # Redraw axes rectangle
                self.plotpanel.canvas.blit(self.axes.bbox)
                
            except RuntimeError:
                print traceback.format_exc()
                print 'len x:', len(self.x_data), 'len y:', len(self.y_data)
                lenx = len(self.x_data)
                leny = len(self.y_data)
                diff = leny - lenx
                if diff > 0:        # if y_data has more items than x_data...
                    self.y_data = self.y_data[:-diff]   # remove items from y_data to even up the arrays
                elif diff < 0:      # if x_data has more items than y_data...
                    self.x_data = self.x_data[:-diff]   # remove items from x_data to even up the arrays
            except:
                print traceback.format_exc()
                    
        self.app.data_lock.release()      # RELEASE DATA_LOCK -->
                
    def target_visible(self, boo):
        self.target.set_visible(boo)
        
    def get_target(self):
        return self.target
    
    def get_axes(self):
        return self.axes
        
    def test_target(self):
        # Test the target line or box and return a boolean value indicating a hit or miss
        result = False
        if self.app.sess_status.mode == 2:          # STRENGTH training
            if self.app.sess_status.peak[1] > self.app.sess_status.str_target:
                result = True

        elif self.app.sess_status.mode == 4:        # SKILL training
            # Test to see if the last peak falls inside the target
            # Start by creating an array containing just the values around the target
            x1 = self.tfactory.x1 - self.app.cfg.target_test_margin
            x2 = self.tfactory.x2 + self.app.cfg.target_test_margin
            if x1 < 0: x1 = 0
            if x2 > self.app.cfg._trial_len: x2 = self.app.cfg._trial_len
            x1 = int((x1*self.app.cfg._sample_rate) + 0.5)
            x2 = int((x2*self.app.cfg._sample_rate) + 0.5)
            ta = self.y_data[x1:x2]
            try:
                x = ta.argmax() + x1        # get the index of the max value inside margins
                y = self.y_data[x]          # record the actual value
                x = x/float(self.app.cfg._sample_rate)
                self.app.sess_status.peak = (x, y)
                if self.tfactory.x1 < x and x < self.tfactory.x2 and self.tfactory.y1 < y and y < self.tfactory.y2:
                    result = True
            except:
                print "error in:", threading.currentThread().name
                print "x1 =", x1
                print "x2 =", x2
                print "tested array =", ta
                print traceback.format_exc()
        return result

    def reset_target(self):
        # Reset the target to baseline values
        if self.app.sess_status.mode == 2:
            self.__draw_target_strength(init=True)
        elif self.app.sess_status.mode == 4:
            self.__draw_target_skill(init=True)
            
        
# End of liveplot.py
